/**
 * Slab 分配器的实现。
 */
#include <cds/list.h>
#include <cds/string.h>
#include <mm/memory.h>
#include <mm/slab.h>
#include <mm/kmalloc.h>
#include <mm/pages.h>
#include <mm/compound_pages.h>
#include <mm/page_alloc.h>
#include <mm/page_flags.h>

static struct kmem_cache kmalloc_cache[KMALLOC_MAX_ORDER] = { 0 };

unsigned long get_kmalloc_free_objs_num(unsigned long order)
{
    return kmalloc_cache[order].free_total;
}

unsigned long get_kmalloc_page_order(unsigned long order)
{
    return kmalloc_cache[order].page_order;
}

/**
 * @brief 
 */
static inline void *__object_idx_to_virtual(
    struct page_struct *page, 
    unsigned char idx
)
{
    struct kmem_cache *s = page->slab_cache;
    if (idx >= s->objs_each) {
        /// idx 越界。
        return NULL;
    }
    return page->first_object + s->size * idx;
}

static inline unsigned char __virtual_to_object_idx(
    struct page_struct *page,
    void *virtual
)
{
    struct kmem_cache *s = page->slab_cache;
    unsigned long idx = (unsigned long)(virtual - page->first_object);
    idx /= s->size;
    if (idx >= s->objs_each) {
        /// idx 越界。
        return 0xff;
    }
    return (unsigned char)idx;
}

/**
 * @brief 将一个 Slab 对象加入其所在的 Slab 描述符。
 * 
 * 该函数会根据空闲对象的计数判断其所在的链表，从中移出，并修改相关计数。
 * 
 * 调用该函数之前，必须将复合页设置为属于 Slab 分配器，并设置其所在的 Slab 描述符。该函数不会
 * 做任何有效性检查，所以必须保证输入的参数有效。
 * 
 * @param page 进行移除的 Slab 对象所在复合页的首页。
 */
static inline void __pull_kmem_object(struct page_struct *page)
{
    struct kmem_cache *s = page->slab_cache;
    if (page->free_objs == 0) {
        s->full_count -= 1;
    } else if (page->free_objs == s->objs_each) {
        s->empty_count -= 1;
    } else {
        s->partial_count -= 1;
    }

    s->free_total -= page->free_objs;
    s->objs_total -= s->objs_each;
    list_del_init(&page->lru);
}
/* 与 __pull_kmem_object 相反的操作。 */
static inline void __put_kmem_object(struct page_struct *page)
{
    struct kmem_cache *s = page->slab_cache;
    if (page->free_objs == 0) {
        s->full_count += 1;
        list_add_tail(&s->full_list, &page->lru);
    } else if (page->free_objs == s->objs_each) {
        s->empty_count += 1;
        list_add_tail(&s->empty_list, &page->lru);
    } else {
        s->partial_count += 1;
        list_add_tail(&s->partial_list, &page->lru);
    }
    s->free_total += page->free_objs;
    s->objs_total += s->objs_each;
}

/**
 * @brief 判断 Slab 对象是否适合使用 KMEM_FLAG_OBJ_FREELIST 方式存放空闲链表。
 * 
 * @param aligned_size 单个内存对象对齐后的大小。
 * @param offset_size 第一个内存对象在复合页中的偏移。
 * @param page_order Slab 对象所在复合页大小的阶数。
 * 
 * @retval true 适合。
 * @retval false 不适合。
 */
static inline int __fit_obj_free_list(
    size_t aligned_size,
    size_t offset_size,
    unsigned long page_order
)
{
    unsigned long size_after_offset = 1UL << (PAGE_SHIFT + page_order);
    size_after_offset -= offset_size;
    unsigned long objs_num = size_after_offset / aligned_size;
    return objs_num <= aligned_size;
}

/**
 * @brief 判断 Slab 对象是否适合使用 KMEM_FLAG_NORMAL 方式存放空闲链表。
 * 
 * @param aligned_size 单个内存对象对齐后的大小。
 * @param offset_size 第一个内存对象在复合页中的偏移。
 * @param page_order Slab 对象所在复合页大小的阶数。
 * 
 * @retval true 适合。
 * @retval false 不适合。
 */
static inline int __fit_normal_free_list(
    size_t aligned_size,
    size_t offset_size,
    unsigned long page_order
)
{
    unsigned long size_after_offset = 1UL << (PAGE_SHIFT + page_order);
    size_after_offset -= offset_size;
    unsigned long objs_num = size_after_offset / aligned_size;
    unsigned long remains = size_after_offset - aligned_size * objs_num;
    return remains >= objs_num;
}

/**
 * @brief 对一个Slab描述符结构体的动态初始化。
 * 
 * 该函数仅仅根据给定的参数将指定的 Slab 描述符初始化，并不做进一步的检查。
 * 
 * 初始化结束后会得到一个空的 Slab 描述符，并且还需要进一步置位需要被置位的状态位。
 * 
 * @param s 进行初始化的 Slab 描述符。
 * @param name  Slab 描述符的名称。
 * @param aligned_size  单个内存对象对齐后的大小。
 * @param alignment     内存对象的对齐大小。
 * @param objs_offset   第一个内存对象在复合页内的偏移。
 * @param objs_each     单个 Slab 对象所容纳的内存对象个数。
 * @param page_order    单个 Slab 对象所在复合页大小的阶数。
 * @param free_limit    整个 Slab 分配器允许存在的空闲对象的上限。
 */
static void __kmem_cache_init(
    struct kmem_cache *s, 
    const char *name, 
    size_t aligned_size, 
    size_t alignment,
    size_t objs_offset,
    unsigned long objs_each,
    unsigned long page_order,
    unsigned long free_limit
) {
    memset(s, 0, sizeof(struct kmem_cache));
    list_head_init(&s->full_list);
    list_head_init(&s->partial_list);
    list_head_init(&s->empty_list);
    init_spinlock_t(&s->lock);
    s->name = (char *)name;
    s->size = aligned_size;
    s->alignment = alignment;
    s->objs_offset = objs_offset;
    s->objs_each = objs_each;
    s->page_order = page_order;
    s->free_limit = free_limit;
}

int kmalloc_init(void)
{
    for (unsigned long i = KMALLOC_MINIMUN_ORDER; i < KMALLOC_MAX_ORDER; i++) {
        /// 计算单个 Slab 对象所需要复合页大小的 order。
        unsigned long page_order = i + 8;
        if (page_order <= PAGE_SHIFT) {
            page_order = 0;
        } else if (page_order >= SLAB_PAGE_MAX_ORDER + PAGE_SHIFT) {
            page_order = SLAB_PAGE_MAX_ORDER - 1;
        } else {
            page_order -= PAGE_SHIFT;
        }

        /// 计算单个 Slab 对象可容纳的内存对象数量。
        size_t aligned_size = 1UL << i; /// 对齐后单个内存对象的大小。
        unsigned long objs_num = 0; /// 单个 Slab 对象可容纳的内存对象数量。
        if (i <= 8) {
            /// 对于单个内存对象小于等于256字节的，强行在最后保留一片区域存放freelist。
            objs_num = (1UL << (PAGE_SHIFT + page_order)) / (aligned_size + 1);
            /// 对于单个对象小于17字节的，需要强行将 objs_each 限制为252
            if (objs_num > 252) {
                objs_num = 252;
            }
            __kmem_cache_init(
                kmalloc_cache + i, 
                "kmalloc_cache", 
                aligned_size, 
                8, 
                0, 
                objs_num, 
                page_order, 
                objs_num
            );
            SetKmemNormalFreelist(kmalloc_cache + i);
            SetKmemStaticCache(kmalloc_cache + i);
            SetKmemValid(kmalloc_cache + i);
        } else {
            /// 对于单个内存对象大于等于256字节的，单独分配一块空间存放freelist。
            objs_num = (1UL << (page_order + PAGE_SHIFT - i)) - 1;
            /// 对于单个对象小于17字节的，需要强行将 objs_each 限制为252
            if (objs_num > 252) {
                objs_num = 252;
            }
            __kmem_cache_init(
                kmalloc_cache + i, 
                "kmalloc_cache", 
                aligned_size, 
                8, 
                0, 
                objs_num, 
                page_order, 
                objs_num
            );
            SetKmemOffFreelist(kmalloc_cache + i);
            SetKmemStaticCache(kmalloc_cache + i);
            SetKmemValid(kmalloc_cache + i);
        }

    }
    return 0;
}

void *kmalloc(size_t size)
{
    if (size > (1UL << (KMALLOC_MAX_ORDER - 1))) {
        return NULL;
    }
    unsigned long order = KMALLOC_MINIMUN_ORDER;
    if (size > (1UL << KMALLOC_MINIMUN_ORDER)) {
        size -= 1;
        size >>= KMALLOC_MINIMUN_ORDER;
        while (size > 0) {
            order += 1;
            size >>= 1;
        }
    }
    return kmem_cache_alloc(kmalloc_cache + order);
}

static unsigned long __kfree_count = 0;

void kfree(void *ptr)
{
    struct page_struct *page = virtual_to_page(ptr);
    if (page == NULL) {
        return;
    }
    if (PageCompound(page)) {
        page = get_header_page(page);
    }
    if (page == NULL) {
        return;
    }
    if (!PageSlab(page)) {
        return;
    }
    kmem_cache_free(page->slab_cache, ptr);
}

void *kmem_cache_alloc(struct kmem_cache *s)
{
    if (!is_kmem_valid(s)) {
        return NULL;
    }
    spinlock_acquire(&s->lock);

    /// 先尝试从部分空闲链表内获得 Slab 对象。
    struct page_struct *page = list_first_entry_or_null(
        &s->partial_list, 
        struct page_struct, 
        lru
    );
    if (page == NULL) {
        /// 如果部分空闲链表中获取失败，尝试从空闲链表中获得。
        page = list_first_entry_or_null(
            &s->empty_list, 
            struct page_struct, 
            lru
        );
        if (page == NULL) {
            /// 如果空闲链表中获取失败，尝试扩展 Slab 分配器。
            page = __kmem_cache_grow(s);
            if (page == NULL) {
                /// 如果扩展失败，返回NULL。
                spinlock_release(&s->lock);
                return NULL;
            }
        } else {
            __pull_kmem_object(page);
        }
    } else {
        __pull_kmem_object(page);
    }

    unsigned char new_obj_idx = page->free_object_entry;
    if (new_obj_idx == 0xff) {
        /// @todo 一个BUG！
while (1);
    }
    page->free_object_entry = *(unsigned char *)(page->free_list + new_obj_idx);
    *(unsigned char *)(page->free_list + new_obj_idx) = 0xff;
    page->free_objs -= 1;
    __put_kmem_object(page);
    spinlock_release(&s->lock);
    return __object_idx_to_virtual(page, new_obj_idx);
}

void kmem_cache_free(struct kmem_cache *s, void *ptr)
{
    if (!is_kmem_valid(s)) {
        return;
    }
    spinlock_acquire(&s->lock);

    struct page_struct *page = virtual_to_page(ptr);
    if (page == NULL) {
        spinlock_release(&s->lock);
        return;
    }
    if (PageCompound(page)) {
        page = get_header_page(page);
        if (page == NULL) {
            spinlock_release(&s->lock);
            return;
        }
    }
    if (!PageSlab(page)) {
        spinlock_release(&s->lock);
        return;
    }
    if (page->slab_cache != s) {
        spinlock_release(&s->lock);
        return;
    }

    unsigned long idx = __virtual_to_object_idx(page, ptr);
    if (idx == 0xff) {
        spinlock_release(&s->lock);
while (1);
        return;
    }
    if (*(unsigned char *)(page->free_list + idx) != 0xff) {
        spinlock_release(&s->lock);
while (1);
        return;
    }

    /// 修改空闲链表。
    __pull_kmem_object(page);
    *(unsigned char *)(page->free_list + idx) = page->free_object_entry;
    page->free_object_entry = idx;
    page->free_objs += 1;
    __put_kmem_object(page);

    /// 尝试让空闲的内存对象不过多。
    if (s->free_total > s->free_limit) {
        __kmem_cache_shrink(s);
    }

    spinlock_release(&s->lock);
}

/**
 * @brief 初始化一个 Slab 对象。
 * 
 * 该函数不会检查输入参数的有效性，仅仅进行初始化。
 */
static inline void __init_slab_object(
    struct kmem_cache *s, 
    struct page_struct *page,
    void *page_virtual, 
    void *slab_freelist
)
{
    /// 初始化 Slab 对象的空闲链表。
    for (unsigned char i = 0; i < s->objs_each; i++) {
        *(unsigned char *)(slab_freelist + i) = i + 1;
    }
    *(unsigned char *)(slab_freelist + s->objs_each - 1) = 0xff;

    /// 在页结构体中保存相关信息。
    page->slab_cache = s;
    page->free_list = slab_freelist;
    page->first_object = page_virtual + s->objs_offset;
    page->free_object_entry = 0;
    page->free_objs = s->objs_each;
    SetPageSlab(page);
}

/**
 * @brief 释放一个 Slab 对象。
 * 
 * 必须保证传入的 Slab 对象已经从其对应的 Slab 描述符中移除。
 * 
 * 该函数不会检查输入参数的有效性，仅仅进行析构。
 */
static inline void __destroy_slab_object(
    struct page_struct *page, 
    unsigned long order
)
{
    struct kmem_cache *s = page->slab_cache;
    if (KmemOffFreelist(s)) {
        /// 对于空闲链表单独分配的 Slab 描述符，需要进行释放。
        kfree(page->free_list);
    }
    ClearPageSlab(page);
    free_pages(page, order);
}

struct page_struct *__kmem_cache_grow(struct kmem_cache *s)
{
    if (!is_kmem_valid(s)) {
        /// 不能为不合法的 Slab 描述符分配新的 Slab 对象。
        return NULL;
    }
    struct page_struct *page = alloc_pages(s->page_order);
    if (page == NULL) {
        /// 页分配失败。
        return NULL;
    }

    void *page_virtual = get_page_virtual(page);
    void *slab_freelist = page_virtual;
    if (KmemOffFreelist(s)) {
        slab_freelist = kmalloc(s->objs_each);
        if (slab_freelist == NULL) {
            /// 为 Slab 对象分配空闲链表空间失败。。
            free_pages(page, s->page_order);
            return NULL;
        }
    } else if (KmemNormalFreelist(s)) {
        slab_freelist = page_virtual;
        slab_freelist -= s->objs_each;
        slab_freelist += 1UL << (PAGE_SHIFT + s->page_order);
    } else if (KmemObjFreelist(s)) {
        slab_freelist = page_virtual;
        slab_freelist += s->objs_offset;
        slab_freelist += s->size * s->objs_each;
    } else {
        /// 未知 Slab 结构。
        free_pages(page, s->page_order);
        return NULL;
    }

    __init_slab_object(s, page, page_virtual, slab_freelist);

    return page;
}

int __kmem_cache_shrink(struct kmem_cache *s)
{
    if (!is_kmem_valid(s)) {
        /// 不能对不合法的 Slab 描述符进行该操作。
        return -1;
    }
    struct page_struct *page = list_first_entry_or_null(
        &s->empty_list, 
        struct page_struct, 
        lru
    );
    if (page == NULL) {
        return -1;
    }
    unsigned long page_order = get_page_order(page);
    if (page_order != s->page_order) {
        /// @todo 一个BUG！
while (1);
        return -1;
    }

    __pull_kmem_object(page);
    __destroy_slab_object(page, page_order);
    
    return 0;
}

void kmem_cache_destroy(struct kmem_cache *s)
{
    if (!is_kmem_valid(s)) {
        /// 不能对不合法的 Slab 描述符进行该操作。
        return;
    }
    if (KmemStaticCache(s)) {
        /// 不能对静态的 Slab 描述符进行该操作。
        return;
    }
    spinlock_acquire(&s->lock);

    struct page_struct *pos;
    struct page_struct *n;
    list_for_each_entry_careful(&s->empty_list, pos, n, lru) {
        __destroy_slab_object(pos, s->page_order);
    }
    list_for_each_entry_careful(&s->partial_list, pos, n, lru) {
        __destroy_slab_object(pos, s->page_order);
    }
    list_for_each_entry_careful(&s->full_list, pos, n, lru) {
        __destroy_slab_object(pos, s->page_order);
    }
    kfree(s);
}

struct kmem_cache *kmem_cache_create(
    const char *name, 
    size_t size, 
    size_t align
)
{
    if (size > SLAB_MAX_OBJ_SIZE) {
        /// 内存区域大小过大。
        return NULL;
    }
    size_t aligned_size = (size + 7) & (~7UL);
    if (aligned_size <= SLAB_MINIMUN_BOJ_SIZE) {
        aligned_size = SLAB_MINIMUN_BOJ_SIZE;
    }

    struct kmem_cache *new_slab = kmalloc(sizeof(struct kmem_cache));
    if (new_slab == NULL) {
        return NULL;
    }

    unsigned long objs_offset = 0;
    size_t tmp_size = aligned_size * 255 + objs_offset;
    size_t order_size = 1UL << PAGE_SHIFT;
    unsigned long order_count = 0;
    unsigned long objs_each = 0;
    while (order_size < tmp_size) {
        order_size <<= 1;
        order_count += 1;
    }
    if (order_count >= SLAB_PAGE_MAX_ORDER) {
        order_count = SLAB_PAGE_MAX_ORDER - 1;
        order_size = 1UL << (SLAB_PAGE_MAX_ORDER + PAGE_SHIFT - 1);
    } else if (order_count > 0) {
        order_count -= 1;
        order_size >>= 1;
    }

    tmp_size = order_size - objs_offset;
    objs_each = tmp_size / aligned_size;
    /// 对于单个对象小于17字节的，需要强行将 objs_each 限制为252
    if (objs_each > 252) {
        objs_each = 252;
    }
    if (__fit_normal_free_list(aligned_size, objs_offset, order_count)) {
        __kmem_cache_init(
            new_slab,
            name,
            aligned_size,
            align,
            objs_offset,
            objs_each,
            order_count,
            objs_each
        );
        SetKmemNormalFreelist(new_slab);
    } else if (__fit_obj_free_list(aligned_size, objs_offset, order_count)) {
        __kmem_cache_init(
            new_slab,
            name,
            aligned_size,
            align,
            objs_offset,
            objs_each - 1,
            order_count,
            objs_each - 1
        );
        SetKmemObjFreelist(new_slab);
    } else {
        __kmem_cache_init(
            new_slab,
            name,
            aligned_size,
            align,
            objs_offset,
            objs_each,
            order_count,
            objs_each
        );
        SetKmemOffFreelist(new_slab);
    }
    SetKmemValid(new_slab);

    return new_slab;
}   

